from keras.layers import Conv1D, Dense, Dropout, BatchNormalization, MaxPooling1D, Activation, Flatten, Input
from keras.models import Sequential
from keras.utils import plot_model
from keras.regularizers import l2
import preprocess
from keras.callbacks import TensorBoard
import numpy as np
import os

# 设置Graphviz路径（重要！）
os.environ["PATH"] += os.pathsep + 'C:/Program Files/Graphviz/bin/'

# 训练参数
batch_size = 128
epochs = 20
num_classes = 10
length = 2048
BatchNorm = True  # 是否批量归一化
number = 1000  # 每类样本的数量
normal = True  # 是否标准化
rate = [0.7, 0.2, 0.1]  # 测试集验证集划分比例

path = r'C:\Users\xsy\Desktop\wdcnn_bearning_fault_diagnosis-master\data\0HP'
x_train, y_train, x_valid, y_valid, x_test, y_test = preprocess.prepro(
    d_path=path, length=length,
    number=number,
    normal=normal,
    rate=rate,
    enc=True, enc_step=28
)

# 增加通道维度
x_train, x_valid, x_test = x_train[:, :, np.newaxis], x_valid[:, :, np.newaxis], x_test[:, :, np.newaxis]
input_shape = x_train.shape[1:]

print('训练样本维度:', x_train.shape)
print(x_train.shape[0], '训练样本个数')
print('验证样本的维度', x_valid.shape)
print(x_valid.shape[0], '验证样本个数')
print('测试样本的维度', x_test.shape)
print(x_test.shape[0], '测试样本个数')


# 定义卷积层
def wdcnn(model, filters, kernerl_size, strides, conv_padding, pool_padding, pool_size, BatchNormal):
    model.add(Conv1D(filters=filters, kernel_size=kernerl_size, strides=strides,
                     padding=conv_padding, kernel_regularizer=l2(1e-4)))
    if BatchNormal:
        model.add(BatchNormalization())
    model.add(Activation('relu'))
    model.add(MaxPooling1D(pool_size=pool_size, padding=pool_padding))
    return model


# 构建模型
model = Sequential()
model.add(Input(shape=input_shape))  # 显式定义Input层（更规范）
model.add(Conv1D(filters=16, kernel_size=64, strides=16, padding='same', kernel_regularizer=l2(1e-4)))
model.add(BatchNormalization())
model.add(Activation('relu'))
model.add(MaxPooling1D(pool_size=2))

# 添加卷积层
model = wdcnn(model, filters=32, kernerl_size=3, strides=1, conv_padding='same',
              pool_padding='valid', pool_size=2, BatchNormal=BatchNorm)
model = wdcnn(model, filters=64, kernerl_size=3, strides=1, conv_padding='same',
              pool_padding='valid', pool_size=2, BatchNormal=BatchNorm)
model = wdcnn(model, filters=64, kernerl_size=3, strides=1, conv_padding='same',
              pool_padding='valid', pool_size=2, BatchNormal=BatchNorm)
model = wdcnn(model, filters=64, kernerl_size=3, strides=1, conv_padding='valid',
              pool_padding='valid', pool_size=2, BatchNormal=BatchNorm)

model.add(Flatten())
model.add(Dense(units=100, activation='relu', kernel_regularizer=l2(1e-4)))
model.add(Dense(units=num_classes, activation='softmax', kernel_regularizer=l2(1e-4)))

# 编译模型
model.compile(optimizer='Adam', loss='categorical_crossentropy', metrics=['accuracy'])

# 训练模型
tb_cb = TensorBoard(log_dir='logs')
model.fit(x=x_train, y=y_train, batch_size=batch_size, epochs=epochs,
          verbose=1, validation_data=(x_valid, y_valid), shuffle=True,
          callbacks=[tb_cb])

# 评估模型
score = model.evaluate(x=x_test, y=y_test, verbose=0)
print("测试集上的损失：", score[0])
print("测试集准确率:", score[1])

# 可视化模型（优化后的代码）
try:
    plot_model(model, to_file='wdcnn_model.png', show_shapes=True, dpi=100, show_layer_names=True)
    print("模型结构图已保存为 wdcnn_model.png")

    # 同时生成文本摘要
    with open('model_summary.txt', 'w') as f:
        model.summary(print_fn=lambda x: f.write(x + '\n'))
    print("模型文本摘要已保存为 model_summary.txt")

except Exception as e:
    print("绘图失败，请检查Graphviz配置:", str(e))
    print("改用文本摘要显示结构：")
    model.summary()